import threading
from .ai_models import TextEmbedding, ImageEmbedding ,TableDetector
from django.conf import settings

# 全局模型实例字典
_models = {}
_lock = threading.Lock()

# 模型类型常量
TEXT_MODEL = "text_embedding"
IMAGE_MODEL = "image_embedding"
TABLE_DETECT_MODEL = "table_detect_model"

def _get_model_key(model_type, model_path):
    """生成唯一的模型"""
    return f"{model_type}_{model_path}"

def get_model(model_type, model_path=None):
    """获取文本嵌入模型实例（线程安全）"""
    global _models

    # 使用默认模型路径（如果未提供）
    if model_path is None:
        if model_type == TEXT_MODEL:
            model_path = settings.TEXT_EMBEDDING_MODEL_PATH
        elif model_type == IMAGE_MODEL:
            model_path = settings.IMAGE_EMBEDDING_MODEL_PATH
        elif model_type == TABLE_DETECT_MODEL:
            model_path = settings.TABLE_DETECT_MODEL_PATH
        else:
            raise ValueError(f"未知模型类型: {model_type}")

    model_key = _get_model_key(model_type, model_path)

    # 双重检查锁定确保线程安全
    if model_key not in _models:
        with _lock:
            if model_key not in _models:
                print(f"🔄 正在加载模型: {model_type} ({model_path})")

                #根据类型创建模型实例
                if model_type == TEXT_MODEL:
                    _models[model_key] = TextEmbedding(model_path)
                elif model_type == IMAGE_MODEL:
                    _models[model_key] = ImageEmbedding(model_path)
                elif model_type == TABLE_DETECT_MODEL:
                    _models[model_key] = TableDetector(model_path=model_path)
                else:
                    raise ValueError(f"不支持模型:{model_type} ({model_path})")
    return _models[model_key]



def preload_models():
    """预加载所有需要的模型"""
    if hasattr(settings, 'PRELOAD_MODELS') and settings.PRELOAD_MODELS:
        print("🚀 预加载模型中...")

        # 预加载文本模型
        if hasattr(settings, 'TEXT_EMBEDDING_MODEL_PATH'):
            get_model(TEXT_MODEL, settings.TEXT_EMBEDDING_MODEL_PATH)

        #预加载图像模型
        if hasattr(settings, 'IMAGE_EMBEDDING_MODEL_PATH'):
            get_model(IMAGE_MODEL, settings.IMAGE_EMBEDDING_MODEL_PATH)

        #预加载表格检测模型
        if hasattr(settings, 'TABLE_DETECT_MODEL_PATH'):
            get_model(TABLE_DETECT_MODEL, settings.TABLE_DETECT_MODEL_PATH)

        print("✅ 所有模型预加载完成")

def get_text_embedding_model(model_path = None):
    """获取文本嵌入模型（快捷方式）"""
    return get_model(TEXT_MODEL, model_path)

def get_image_embedding_model(model_path = None):
    """获取图像嵌入模型（快捷方式）"""
    return get_model(IMAGE_MODEL, model_path)

def get_table_detect_model(model_path = None):
    """获取表格检测模型（快捷方式）"""
    return get_model(TABLE_DETECT_MODEL, model_path)